avgGradients = [];
avgSquaredGradients = [];
doTraining = true;

if doTraining
    
    % Create a minibatchqueue to batch data from training and validation
    % datastores. Use the batchData function, listed at the end of the
    % example, to batch the point cloud data and one-hot encode the label 
    % data.
    numOutputsFromDSRead = 2;
    mbqTrain = minibatchqueue(dsTrain,numOutputsFromDSRead,...
        "MiniBatchSize", miniBatchSize,...
        "MiniBatchFcn",@batchData,...
        "MiniBatchFormat",["SCSB" "BC"]);
    
    mbqVal = minibatchqueue(dsVal,numOutputsFromDSRead,...
        "MiniBatchSize", miniBatchSize,... 
        "MiniBatchFcn",@batchData,...
        "MiniBatchFormat",["SCSB" "BC"]);
 
    % Use the configureTrainingProgressPlot function, listed at the end of the
    % example, to initialize the training progress plot to display the training
    % loss, training accuracy, and validation accuracy.
    [lossPlotter, trainAccPlotter,valAccPlotter] = initializeTrainingProgressPlot;
    
    numClasses = numel(classes);
    iteration = 0;
    start = tic;
    for epoch = 1:numEpochs
        
        % Shuffle data every epoch.
        shuffle(mbqTrain);
      
        % Iterate through data set.
        while hasdata(mbqTrain)
            iteration = iteration + 1;
            
            % Read next batch of training data.
            [XTrain, YTrain] = next(mbqTrain);            
            
            % Evaluate the model gradients and loss using dlfeval and the
            % modelGradients function.
            [gradients, loss, state, acc] = dlfeval(@modelGradients,XTrain,YTrain,parameters,state);
            
            % L2 regularization.
            gradients = dlupdate(@(g,p) g + l2Regularization*p,gradients,parameters);
            
            % Update the network parameters using the Adam optimizer.
            [parameters, avgGradients, avgSquaredGradients] = adamupdate(parameters, gradients, ...
                avgGradients, avgSquaredGradients, iteration,...
                learnRate,gradientDecayFactor, squaredGradientDecayFactor);
            
            % Update the training progress.
            D = duration(0,0,toc(start),"Format","hh:mm:ss");
            title(lossPlotter.Parent,"Epoch: " + epoch + ", Elapsed: " + string(D))
            addpoints(lossPlotter,iteration,double(gather(extractdata(loss))))
            addpoints(trainAccPlotter,iteration,acc);
            drawnow
        end
        
        % Evaluate the model on validation data.
        cmat = sparse(numClasses,numClasses);
        while hasdata(mbqVal)
            
            % Read next batch of validation data.
            [XVal, YVal] = next(mbqVal);

            % Compute label predictions.
            isTraining = false;
            YPred = pointnetClassifier(XVal,parameters,state,isTraining);
            
            % Choose prediction with highest score as the class label for
            % XTest.
            [~,YValLabel] = max(YVal,[],1);
            [~,YPredLabel] = max(YPred,[],1);
            
            % Collect confusion metrics.
            cmat = aggreateConfusionMetric(cmat,YValLabel,YPredLabel);
        end
        
        % Update training progress plot with average classification accuracy.
        acc = sum(diag(cmat))./sum(cmat,"all");
        addpoints(valAccPlotter,iteration,acc);
        
        % Upate the learning rate.
        if mod(epoch,learnRateDropPeriod) == 0
            learnRate = learnRate * learnRateDropFactor;
        end
        
        % Reset training and validation data queues.
        reset(mbqTrain);
        reset(mbqVal);
    end

else
    % Download pretrained model parameters, model state, and validation
    % results.
    pretrainedURL = 'https://ssd.mathworks.com/supportfiles/vision/data/pointnetSydneyUrbanObjects.zip'; 
    pretrainedResults = downloadPretrainedPointNet(pretrainedURL);
   
    parameters = pretrainedResults.parameters;
    state = pretrainedResults.state;
    cmat = pretrainedResults.cmat;
    
    % Move model parameters to the GPU if possible and convert to a dlarray.
    parameters = prepareForPrediction(parameters,@(x)dlarray(toDevice(x,canUseGPU)));
    
    % Move model state to the GPU if possible.
    state = prepareForPrediction(state,@(x)toDevice(x,canUseGPU));
end
